#include <vector>
#include <map>
using namespace std;

#include "hla_processor.h"

static void get_haplotype_peptides(const Haplotype& haplotype, set<string>& peptides);

class AveragesProcessor : public HlaGroupsPairProcessor
{
public:
    AveragesProcessor() {}
private:
    bool enabled;

    typedef vector<float> HaplotypeData;

    map<string, HaplotypeData> haplotypes;

    void generate_headers(Sheet& sheet) const;
    void generate_sheet(Spreadsheet& spreadsheet, const string& haplotype, const HaplotypeData& haplotypeData) const;

    void set_groups(const string& group1, const string& group2)
    {
        enabled = (group1 != "SIV" && group2 != "SIV");
    }

    void add_haplotype(const string& haplotype, size_t only1, size_t common, size_t only2)
    {
        if (enabled)
            haplotypes[haplotype].push_back(float(common) / float(only1+common+only2));
    }

    virtual ProcessorPriority get_priority() const  {   return 1.5f;    }
    virtual void generate_results(Spreadsheet& spreadsheet);
    virtual void process(const Group& g1, const Group& g2);

    void compare_haplotypes(const Haplotype& haplotype1, const Haplotype& haplotype2);
};

REGISTER_PROCESSOR(AveragesProcessor);


template <class T>
static void check_all_in(const T& container2, const T& container1)
{
    // this is O(n^2), could be improved with a set.

    for (typename T::const_iterator it = container2.begin(); it != container2.end(); ++it)
    {
        if (find(container1.begin(), container1.end(), *it) == container1.end())
            cerr << "ERROR: Element " << it->name << " not found in counterpart!" << endl;
    }
}


template <class T>
static size_t sets_difference_count(const set<T>& set1, const set<T>& set2)
{
    set<T> only1;

    set_difference(set1.begin(), set1.end(), set2.begin(), set2.end(), inserter(only1, only1.begin()));

    return only1.size();
}

template <class T>
static size_t sets_common_count(const set<T>& set1, const set<T>& set2)
{
    set<T> common;

    set_intersection(set1.begin(), set1.end(), set2.begin(), set2.end(), inserter(common, common.begin()));

    return common.size();
}

void AveragesProcessor::process(const Group& group1, const Group& group2)
{
    set_groups(group1.name, group2.name);

    // for each haplotype in group1, search it's equivalent in group2, and process:
    for (list<Haplotype>::const_iterator haplotype1 = group1.haplotypes.begin(); haplotype1 != group1.haplotypes.end(); ++haplotype1)
    {
        // search equivalent haplotype in group2:
        // TODO: REIMPLEMENT USING std::map (needs FXP support)
        const list<Haplotype>::const_iterator haplotype2 = find(group2.haplotypes.begin(), group2.haplotypes.end(), *haplotype1);

        if (haplotype2 != group2.haplotypes.end())
            compare_haplotypes(*haplotype1, *haplotype2);
        else
            cerr << "ERROR: Haplotype " << haplotype1->name << " not found in group " << group2.name << endl;
    }

    check_all_in(group2.haplotypes, group1.haplotypes);
}

void AveragesProcessor::compare_haplotypes(const Haplotype& haplotype1, const Haplotype& haplotype2)
{
    set<string> peptides1;
    set<string> peptides2;

    // for each sequence in haplotype1, get the set of peptides. Same for haplotype2:
    get_haplotype_peptides(haplotype1, peptides1);
    get_haplotype_peptides(haplotype2, peptides2);

    add_haplotype(haplotype1.name,
        sets_difference_count(peptides1, peptides2),
        sets_common_count(peptides1, peptides2),
        sets_difference_count(peptides2, peptides1)
    );
}

static void get_haplotype_peptides(const Haplotype& haplotype, set<string>& peptides)
{
    for (list<Sequence>::const_iterator sequence = haplotype.sequences.begin(); sequence != haplotype.sequences.end(); ++sequence)
        peptides.insert(sequence->peptides.begin(), sequence->peptides.end());
}

void AveragesProcessor::generate_headers(Sheet& sheet) const
{
    Row header(sheet);
    const Style center( Style::ALIGN_CENTER | Style::BORDER_BOTTOM );

    for (map<string, HaplotypeData>::const_iterator hla = haplotypes.begin();
         hla != haplotypes.end();
         ++hla)
    {
        header << center << hla->first << v_separator;
    }
}

void AveragesProcessor::generate_results(Spreadsheet& spreadsheet)
{
    Sheet new_sheet( spreadsheet, "Averages" );
    generate_headers(new_sheet);

    const size_t columns = haplotypes.size();

    if (columns > 0)
    {
        HaplotypeData::const_iterator* const columns_it = new HaplotypeData::const_iterator[columns];
        HaplotypeData::const_iterator* const ends       = new HaplotypeData::const_iterator[columns];

        size_t column(0);

        for (map<string, HaplotypeData>::const_iterator hla = haplotypes.begin();
             hla != haplotypes.end();
             ++hla)
        {
            columns_it[column] = hla->second.begin();
            ends[column]       = hla->second.end();
            column++;
        }

        while (columns_it[0] != ends[0])
        {
            Row data(new_sheet);
            for (column = 0; column < columns; ++column)
            {
                data << *(columns_it[column]) << v_separator;
                ++columns_it[column];
            }
        }

        delete [] columns_it;
        delete [] ends;
    }
}

